import numpy as np
import queue
import matplotlib.pyplot as plt
from matplotlib.patches import ConnectionPatch, Circle

def findShortestPath(adjacency_matrix, start, end):
    maxLen = 10000
    node_num = len(adjacency_matrix)

    path_len = [maxLen for _ in range(node_num)]
    node_flag = [False for _ in range(node_num)]
    up_node_info = {}
    priority_queue_instance = queue.PriorityQueue()
    path_len[start] = 0
    priority_queue_instance.put((0, start))

    while not priority_queue_instance.empty():
        current_len, current_node = priority_queue_instance.get()
        if node_flag[current_node]:
            continue
        node_flag[current_node] = True

        for i in range(node_num):
            if adjacency_matrix[current_node][i] > 0 and not node_flag[i]:
                new_len = current_len + adjacency_matrix[current_node][i]
                if new_len < path_len[i]:
                    path_len[i] = new_len
                    up_node_info[i] = current_node
                    priority_queue_instance.put((new_len, i))

    path_info = get_shortest_path(up_node_info, start, node_num)
    return path_info[end], path_len[end]


def get_shortest_path(up_node_info, start, node_num):
    path_info = {}
    for i in range(node_num):
        if i != start:
            path = []
            path_info[i] = get_path_process(up_node_info, start, i, path)
    return path_info


def get_path_process(up_node_info, start_node, end_node, path):
    if end_node in up_node_info:
        path.insert(0, end_node)
        if up_node_info[end_node] == start_node:
            path.insert(0, start_node)
            return path
        else:
            return get_path_process(up_node_info, start_node, up_node_info[end_node], path)
    return []


def visualize_graph(adjacency_matrix, path, start, end):
    node_num = len(adjacency_matrix)
    pos = {i: (np.cos(2 * np.pi * i / node_num), np.sin(2 * np.pi * i / node_num)) for i in range(node_num)}

    fig, ax = plt.subplots()
    for i in range(node_num):
        for j in range(node_num):
            if adjacency_matrix[i][j] > 0:
                line = ConnectionPatch(pos[i], pos[j], "data", "data", arrowstyle="<-", mutation_scale=20, lw=0.5, color="gray")
                ax.add_patch(line)
                ax.text((pos[i][0]+pos[j][0])/2, (pos[i][1]+pos[j][1])/2, adjacency_matrix[i][j], fontsize=8, ha='center', va='center')

    # 绘制节点
    for i in range(node_num):
        circle = Circle(pos[i], 0.05, edgecolor='black', facecolor='lightblue')
        ax.add_patch(circle)
        ax.text(pos[i][0], pos[i][1], str(i), ha='center', va='center', fontsize=12)

    # 绘制最短路径
    for i in range(len(path) - 1):
        line = ConnectionPatch(pos[path[i]], pos[path[i+1]], "data", "data", arrowstyle="-|>", mutation_scale=20, lw=2, color="red")
        ax.add_patch(line)

    ax.scatter([], [], c='red', label=f'Shortest Path from {start} to {end}')
    ax.legend(loc='upper right')
    ax.set_xlim(-1.5, 1.5)
    ax.set_ylim(-1.5, 1.5)
    ax.axis('equal')
    ax.axis('off')
    plt.title('Graph Visualization with Shortest Path')
    plt.show()


if __name__ == '__main__':
    node_num = 6
    adjacency_matrix = np.zeros((node_num, node_num))
    edges = [
        (0, 1, 3), (0, 5, 4),
        (1, 2, 2), (1, 4, 4), (1, 5, 1),
        (2, 3, 7),
        (4, 2, 3), (4, 3, 2),
        (5, 4, 2)
    ]
    for u, v, w in edges:
        adjacency_matrix[u, v] = w

    print("Adjacency Matrix:")
    print(adjacency_matrix)

    path, length = findShortestPath(adjacency_matrix, 0, 4)
    print(f"Shortest path from 0 to 4: {' -> '.join(map(str, path))}, Length: {length}")

    visualize_graph(adjacency_matrix, path, 0, 4)